#!/Pub/Users/wangyk/miniconda3/envs/pyscenic/bin/python3.10
import logging
import sys
import os
import argparse
import subprocess


def activate_conda_env(env_name):
    # 构建激活环境的命令
    command = f"source /Pub/Users/wangyk/miniconda3/bin/activate {env_name} && env"

    # 使用 subprocess.run 执行命令并捕获输出
    result = subprocess.run(command, capture_output=True,
                            text=True, shell=True, executable="/bin/bash")

    if result.returncode != 0:
        print(f"Error activating conda environment: {result.stderr}")
        return

    # 解析环境变量
    env_vars = result.stdout.splitlines()
    for var in env_vars:
        key, value = var.split("=", 1)
        os.environ[key] = value

activate_conda_env("/Pub/Users/wangyk/miniconda3/envs/pyscenic")


wrp = '''
利用pyscenic进行转录因子识别，只需要seurat对象路径以及结果输出路径。默认为人类数据集，如果是小鼠的单细胞数据集，需要替换三个文件。使用参考：
pySCENIC \
-s /pub/user/your_name/your_path/seurat_obj.rds \
-od /pub/user/your_name/your_path/sceinc_out/
'''

parser = argparse.ArgumentParser(prog='pySCENIC',
                                 description=wrp,
                                 formatter_class=argparse.RawTextHelpFormatter)

parser.add_argument('-s', '--seurat_obj', help="seurat对象路径，rds或者rdata",
                    default="/Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/out/SC_part/05.macro_sub_cellchat/ss_m_sub_anno_ready.rds", metavar='')
parser.add_argument('-od', '--output_dir', help="结果输出路径",
                    default='/Pub/Users/wangyk/project/Poroject/F240510002_OV_marco/test/', metavar='')
parser.add_argument('-t', '--type',
                    help="分析所用数据类型，貌似seurat中的counts或者data都可以。默认data",
                    default='data',
                    metavar='')
parser.add_argument('-tfs', '--tfs',
                    help="有默认值，一般不用修改。参考：https://github.com/aertslab/pySCENIC/tree/master/resources",
                    default="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/hs_hgnc_tfs.txt",
                    metavar='')
parser.add_argument('-f', '--feather',
                    help="有默认值，不用修改，参考：https://resources.aertslab.org/cistarget/databases/homo_sapiens/hg19/refseq_r45/mc9nr/gene_based/hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather",
                    default="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/hg19-tss-centered-10kb-10species.mc9nr.genes_vs_motifs.rankings.feather",
                    metavar='')
parser.add_argument('-tbl', '--tbl',
                    help="有默认值，不用修改，参考：https://resources.aertslab.org/cistarget/motif2tf/motifs-v9-nr.hgnc-m0.001-o0.0.tbl",
                    default="/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/out/4.sc_houxu/十.3/scenic_data/motifs-v9-nr.hgnc-m0.001-o0.0.tbl",
                    metavar='')
parser.add_argument('-o', '--organism',
                    help="有默认值，hsa，如果是鼠类数据，可更改为mmu",
                    default="hsa",
                    metavar='')

args = parser.parse_args()


tfs = args.tfs
feather = args.feather
tbl = args.tbl

if args.organism == 'mmu':
    tfs = args.tfs
    feather = args.feather
    tbl = args.tbl

base_files = {tfs, feather, tbl}
od = args.output_dir
input_loom = f"{args.output_dir}/seurat_obj_expression.loom"


def get_logger(name:str):
    """
    输入logger名称即可。
    """
    formatter = logging.Formatter('%(asctime)s-%(levelname)s -%(message)s')
    # 创建一个logger
    logger = logging.getLogger(name)
    logger.setLevel(logging.INFO)

    # 创建一个控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    # 创建一个文件处理器
    console_handler.setFormatter(formatter)
    # 将处理器添加到logger
    logger.addHandler(console_handler)

    return logger


logger = get_logger(name='pySCENIC')


def ensure_directory_exists(directory_path: str):
    "检查文件路径是否存在"
    if not os.path.exists(directory_path):
        # logger.info(f"递归创建结果路径：{args.output_dir}")
        os.makedirs(directory_path)


def file_check(file_path: str):
    "检查文件是否存在"
    if not os.path.exists(file_path):
        logger.error(f"{file_path} not exists")
        sys.exit(1)
    else:
        pass


def main():
    ensure_directory_exists(args.output_dir)

    for i in base_files:
        file_check(i)

    r_seurat2loom_exprMat_path = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/py/SingleCell/utils4cmd/pysceinc/seurat2loom_exprMat.r"

    logger.info("  1.0 convert seurat_exprMat to loom...")
    r_seurat2loom_exprMat_path_command = [
        r_seurat2loom_exprMat_path,
        args.seurat_obj,
        args.output_dir,
        '-t', args.type
    ]
    result = subprocess.run(r_seurat2loom_exprMat_path_command, capture_output=True, text=True)

    logger.info("  2.1 pyscenic grn...")
    grn_command = [
        "pyscenic", "grn",
        "--num_workers", "36",
        "--output", f"{od}/adj.sample.tsv",
        "--method", "grnboost2",
        input_loom,
        tfs
    ]
    result = subprocess.run(grn_command, capture_output=True, text=True)

    logger.info("  2.2 pyscenic ctx...")
    cistarget_command = [
        "pyscenic", "ctx",
        f"{od}/adj.sample.tsv", feather,
        "--annotations_fname", tbl,
        "--expression_mtx_fname", input_loom,
        "--mode", "dask_multiprocessing",
        "--output", f"{od}/reg.csv",
        "--num_workers", "36",
        "--mask_dropouts"
    ]
    result = subprocess.run(cistarget_command, capture_output=True, text=True)

    logger.info("  2.3 pyscenic aucell...")
    aucell_command = [
        "pyscenic", "aucell", input_loom, f"{od}/reg.csv",
        "--output", f"{od}/pySCENIC_res.loom",
        "--num_workers", "36"
    ]

    result = subprocess.run(aucell_command, capture_output=True, text=True)

    logger.info("  2.4 Convert pvSCENIC result loom file to rds file ...")
    r_sceinc_res_loom2rds_path = "/Pub/Users/wangyk/Project_wangyk/Codelib_YK/codeLibrary_yk/py/SingleCell/utils4cmd/pysceinc/sceinc_res_loom2rds.r"
    os.system(
        f"{r_sceinc_res_loom2rds_path} {od}/pySCENIC_res.loom")

    logger.info("  2.5 分析完成 ")


if __name__ == '__main__':
    main()

